import numpy as np
import deepxde as dde
import deepxde.backend as bkd
from AI4XDE.cases import InverseCase
from AI4XDE.utils.Visualization import Visualization_2D


class Wave_1D(InverseCase):
    def __init__(
        self,
        c,
        interval=[0, 10],
        time_interval=[0, 10],
        NumDomain=0,
        loss_weights=[0.1, 1],
        layer_size=[2] + [25] * 7 + [1],
        activation="silu",
        initializer="Glorot normal",
    ):
        self.c = c
        self.interval = interval
        self.time_interval = time_interval

        self.c_trainable = dde.Variable(c + self.rand_num_between(-1, 1))
        super().__init__(
            "Wave equation",
            external_trainable_variables=[self.c_trainable],
            NumDomain=NumDomain,
            use_output_transform=False,
            layer_size=layer_size,
            activation=activation,
            initializer=initializer,
            metrics=None,
            loss_weights=loss_weights,
            visualization=Visualization_2D(
                x_limit=time_interval,
                y_limit=interval,
                x_label="t",
                y_label="x",
                feature_transform=lambda X: X[:, [1, 0]],
            ),
        )

    def rand_num_between(self, l, h):
        return l + np.random.rand() * (h - l)

    def sol(self, X):
        x = X[:, 0:1]
        t = X[:, 1:2]
        return np.cos(2 * np.pi * (t - x / self.c))

    def get_observe_data(self):
        X = self.geomtime.uniform_points(300 * 300)
        X_rand_choice = np.random.choice(X.shape[0], 4000, replace=False)
        X = X[X_rand_choice, :]
        y = self.sol(X)
        return X, y

    def gen_data(self):
        observe_x, observe_y = self.get_observe_data()
        real_data = dde.icbc.PointSetBC(
            observe_x, observe_y, component=0, batch_size=1000
        )
        data = dde.data.TimePDE(
            self.geomtime,
            self.pde,
            [real_data],
            num_domain=self.NumDomain,
            train_distribution="pseudo",
            solution=self.sol,
            num_test=10000,
        )
        data.add_anchors(observe_x)
        return data

    def gen_geomtime(self):
        geom = dde.geometry.Interval(self.interval[0], self.interval[1])
        timedomain = dde.geometry.TimeDomain(
            self.time_interval[0], self.time_interval[1]
        )
        return dde.geometry.GeometryXTime(geom, timedomain)

    def gen_pde(self):
        def pde(x, y):
            dy_tt = dde.grad.hessian(y, x, i=1, j=1)
            dy_xx = dde.grad.hessian(y, x, i=0, j=0)
            return dy_tt - self.c_trainable**2 * dy_xx

        return pde

    def plot_result(self, solver, colorbar=[0, 0, 0]):
        X = np.array(
            [
                [x, t]
                for x in np.linspace(self.interval[0], self.interval[1], 1000)
                for t in np.linspace(self.time_interval[0], self.time_interval[1], 1000)
            ]
        )
        y = self.sol(X)
        model_y = solver.model.predict(X)

        fig, axes = self.Visualization.plot_exact_predict_error_2D(
            X,
            y,
            model_y,
            shape=[1000, 1000],
            title=f"{solver.name} {self.c}",
            colorbar=colorbar,
        )

        C1_true = self.c
        C1_pred = bkd.to_numpy(self.c_trainable)
        C1_error = np.abs(C1_true - C1_pred)
        print(f"C true: {C1_true}")
        print(f"C pred: {C1_pred}")
        print(f"C error: {C1_error}")

        return fig, axes
